#!/usr/bin/env python3
"""
完整流程测试脚本
测试从注册到启动命令接收的完整流程
"""

import requests
import json
import time
import threading
from typing import Dict, Any

# 配置
CONTROLLER_URL = "http://localhost:8000"
NODE_MANAGER_URL = "http://localhost:8080"


def test_register_and_start_flow():
    """测试完整的注册和启动流程"""
    
    print("=== 完整流程测试 ===")
    print("1. 发送注册请求到 Controller")
    print("2. 等待 Controller 发送启动命令到 NodeManager")
    print("3. 验证启动命令接收")
    print()
    
    # 步骤1: 发送注册请求
    print("步骤1: 发送注册请求...")
    
    register_data = {
        "job_name": "test_flow_job",
        "role": "worker",
        "pod_ip": "10.0.0.1",
        "host_ip": "192.168.1.100",
        "bussiness_port": ["8080"],
        "mgmt_port": "9090",
        "parallel_config": {
            "dp_size": 1,
            "cp_size": 1,
            "tp_size": 1,
            "sp_size": 1,
            "ep_size": 1,
            "pp_size": 1,
            "world_size": 1
        },
        "ranktable": {
            "version": "1.0",
            "status": "active",
            "server_count": "1",
            "server_list": [
                {
                    "server_id": "server_0",
                    "host_ip": "192.168.1.100",
                    "device": [
                        {
                            "device_id": "0",
                            "device_ip": "192.168.1.100",
                            "super_device_id": None,
                            "rank_id": "0"
                        }
                    ]
                }
            ]
        }
    }
    
    try:
        # 发送注册请求
        response = requests.post(f"{CONTROLLER_URL}/controller/register", json=register_data, timeout=30)
        
        if response.status_code == 200:
            result = response.json()
            print(f"✅ 注册成功: {json.dumps(result, indent=2, ensure_ascii=False)}")
        else:
            print(f"❌ 注册失败: {response.status_code} - {response.text}")
            return False
            
    except Exception as e:
        print(f"❌ 注册请求失败: {e}")
        return False
    
    # 步骤2: 等待启动命令
    print("\n步骤2: 等待启动命令...")
    print("请确保 NodeManager mock 服务器正在运行 (python simple_start_receiver.py)")
    print("等待 Controller 发送启动命令...")
    
    # 等待一段时间让 Controller 处理注册并发送启动命令
    time.sleep(5)
    
    # 步骤3: 检查是否收到启动命令
    print("\n步骤3: 检查启动命令接收情况...")
    
    try:
        # 检查 NodeManager 状态
        status_response = requests.get(f"{NODE_MANAGER_URL}/node-manager/status", timeout=10)
        if status_response.status_code == 200:
            status = status_response.json()
            print(f"NodeManager 状态: {json.dumps(status, indent=2, ensure_ascii=False)}")
            
            if status.get("received_commands", 0) > 0:
                print("✅ 成功接收到启动命令!")
                
                # 获取详细命令信息
                commands_response = requests.get(f"{NODE_MANAGER_URL}/node-manager/commands", timeout=10)
                if commands_response.status_code == 200:
                    commands = commands_response.json()
                    print(f"接收到的命令详情: {json.dumps(commands, indent=2, ensure_ascii=False)}")
                return True
            else:
                print("❌ 未收到启动命令")
                return False
        else:
            print(f"❌ 无法连接到 NodeManager: {status_response.status_code}")
            return False
            
    except Exception as e:
        print(f"❌ 检查启动命令时发生错误: {e}")
        return False


def test_multiple_registrations():
    """测试多个注册的流程"""
    
    print("\n=== 多注册测试 ===")
    
    jobs = [
        {
            "job_name": "multi_job_1",
            "role": "worker",
            "pod_ip": "10.0.0.2",
            "host_ip": "192.168.1.101"
        },
        {
            "job_name": "multi_job_2", 
            "role": "coordinator",
            "pod_ip": "10.0.0.3",
            "host_ip": "192.168.1.102"
        }
    ]
    
    for i, job in enumerate(jobs, 1):
        print(f"\n注册作业 {i}: {job['job_name']}")
        
        register_data = {
            "job_name": job["job_name"],
            "role": job["role"],
            "pod_ip": job["pod_ip"],
            "host_ip": job["host_ip"],
            "bussiness_port": ["8080"],
            "mgmt_port": "9090",
            "parallel_config": {
                "dp_size": 1,
                "cp_size": 1,
                "tp_size": 1,
                "sp_size": 1,
                "ep_size": 1,
                "pp_size": 1,
                "world_size": 1
            },
            "ranktable": {
                "version": "1.0",
                "status": "active",
                "server_count": "1",
                "server_list": [
                    {
                        "server_id": f"server_{i}",
                        "host_ip": job["host_ip"],
                        "device": [
                            {
                                "device_id": "0",
                                "device_ip": job["host_ip"],
                                "super_device_id": None,
                                "rank_id": "0"
                            }
                        ]
                    }
                ]
            }
        }
        
        try:
            response = requests.post(f"{CONTROLLER_URL}/controller/register", json=register_data, timeout=30)
            if response.status_code == 200:
                print(f"✅ {job['job_name']} 注册成功")
            else:
                print(f"❌ {job['job_name']} 注册失败: {response.status_code}")
        except Exception as e:
            print(f"❌ {job['job_name']} 注册失败: {e}")
        
        time.sleep(2)  # 避免请求过快
    
    # 等待所有启动命令
    print("\n等待所有启动命令...")
    time.sleep(10)
    
    # 检查所有命令
    try:
        commands_response = requests.get(f"{NODE_MANAGER_URL}/commands", timeout=10)
        if commands_response.status_code == 200:
            commands = commands_response.json()
            print(f"总共接收到 {commands['total_commands']} 个启动命令")
            for cmd in commands['commands']:
                print(f"  - {cmd['job_name']} (实例ID: {cmd['instance_id']})")
    except Exception as e:
        print(f"检查命令时发生错误: {e}")


def main():
    """主函数"""
    print("=== 完整流程测试脚本 ===")
    print("这个脚本测试从注册到启动命令接收的完整流程")
    print()
    print("前置条件:")
    print("1. Controller 服务正在运行 (端口 8000)")
    print("2. NodeManager mock 服务器正在运行 (端口 8080)")
    print()
    
    # 检查服务是否可用
    print("检查服务可用性...")
    
    try:
        # 检查 Controller
        controller_response = requests.get(f"{CONTROLLER_URL}/", timeout=5)
        if controller_response.status_code == 200:
            print("✅ Controller 服务可用")
        else:
            print("❌ Controller 服务不可用")
            return
    except:
        print("❌ 无法连接到 Controller 服务")
        print("请确保 Controller 正在运行在端口 8000")
        return
    
    try:
        # 检查 NodeManager
        nm_response = requests.get(f"{NODE_MANAGER_URL}/", timeout=5)
        if nm_response.status_code == 200:
            print("✅ NodeManager mock 服务可用")
        else:
            print("❌ NodeManager mock 服务不可用")
            return
    except:
        print("❌ 无法连接到 NodeManager mock 服务")
        print("请运行: python simple_start_receiver.py")
        return
    
    print()
    
    # 运行测试
    print("开始测试...")
    
    # 测试1: 单个注册
    success1 = test_register_and_start_flow()
    
    if success1:
        print("\n✅ 单个注册测试通过!")
    else:
        print("\n❌ 单个注册测试失败!")
    
    # 测试2: 多个注册
    test_multiple_registrations()
    
    print("\n=== 测试完成 ===")
    print("请检查 NodeManager mock 服务器日志查看接收到的启动命令详情")


if __name__ == "__main__":
    main()
